import pandas as pd
import matplotlib.pyplot as plt
from collections import Counter
import os

# --- 配置区 ---
# 请将这里的路径替换为您自己的文件路径
INPUT_FILE_1 = 'outputs_100/throughputs_1_10.txt'
INPUT_FILE_2 = 'outputs_100/throughputs_2_10.txt'

# INPUT_FILE_1 = 'outputs_100/throughputs_0_9.txt'
# INPUT_FILE_2 = 'outputs_100/throughputs_3_10.txt'

INTERMEDIATE_FILE_1 = 'inter_high_throughputs.csv'
INTERMEDIATE_FILE_2 = 'inter_low_throughputs.csv'

OUTPUT_PLOT_FILE = 'throughput_comparison.svg'

# --- 脚本主体 ---

def process_file_to_intermediate(input_file, output_file):
    """
    读取日志文件，按5行分组处理，计算真实的吞吐量(tokens/sec)，
    并生成包含聚合数据的中间文件。
    """
    print(f"正在处理文件: {input_file} -> {output_file}")
    
    column_names = ['timestamp', 'dp_rank', 'batch_size', 'draft_tokens', 'time_cost', 'num_generate_tokens']
    
    try:
        df = pd.read_csv(
            input_file, 
            header=None,
            skiprows=1,
            names=column_names
        )
    except FileNotFoundError:
        print(f"错误: 输入文件 '{input_file}' 未找到。")
        return
    except pd.errors.EmptyDataError:
        print(f"警告: 输入文件 '{input_file}' 在跳过表头后为空。")
        return

    if df.empty:
        print(f"警告: 输入文件 '{input_file}' 读取后数据为空。")
        return

    first_timestamp = df['timestamp'].iloc[0]
    
    processed_data = []
    
    for i in range(0, len(df), 5):
        chunk = df.iloc[i:i+5]
        
        if chunk.empty:
            continue
        
        # --- 修改点 1: 正确计算吞吐量 ---
        # 计算5行内的平均生成token数和平均时间成本
        avg_tokens = chunk['num_generate_tokens'].mean()
        avg_time_cost = chunk['time_cost'].mean()
        
        # 计算吞吐量 (tokens/second)，并处理 avg_time_cost 可能为0的情况
        if avg_time_cost > 0:
            throughput_tps = avg_tokens / avg_time_cost
        else:
            throughput_tps = 0.0 # 如果时间为0，吞吐量也记为0
            print(f"警告: 在文件 {input_file} 的行 {i+2}-{i+6} 中, 平均 time_cost 为 0，吞吐量计为 0。")

        if not chunk['batch_size'].empty:
            batch_size_counts = Counter(chunk['batch_size'])
            most_common_bs = batch_size_counts.most_common(1)[0][0]
        else:
            most_common_bs = None

        chunk_start_timestamp = chunk['timestamp'].iloc[0]
        relative_time = chunk_start_timestamp - first_timestamp
        
        # --- 修改点 2: 使用新的吞吐量值和列名 ---
        processed_data.append({
            'relative_time_s': relative_time,
            'throughput_tokens_per_sec': throughput_tps, # <--- 新的列名和值
            'batch_size': most_common_bs
        })
        
    if processed_data:
        intermediate_df = pd.DataFrame(processed_data)
        intermediate_df.to_csv(output_file, index=False)
        print(f"中间文件 '{output_file}' 已成功生成。")
    else:
        print(f"文件 '{input_file}' 中没有足够的数据进行处理。")


def plot_from_intermediate_files(file1, file2, label1, label2, output_plot):
    """
    从两个中间文件中读取数据，并绘制对比曲线图。
    """
    print(f"正在从 '{file1}' 和 '{file2}' 生成图表...")
    
    try:
        df1 = pd.read_csv(file1)
        df2 = pd.read_csv(file2)
    except FileNotFoundError as e:
        print(f"错误: 无法读取中间文件 '{e.filename}'。请先确保已生成该文件。")
        return

    # 检查绘图所需列是否存在
    required_col = 'throughput_tokens_per_sec'
    if required_col not in df1.columns or required_col not in df2.columns:
        print(f"错误: 中间文件缺少 '{required_col}' 列。请确认处理脚本已更新。")
        return

    fig, ax = plt.subplots(figsize=(9.5, 5))
    
    ax.grid(True, alpha=0.3)  # 设置网格线透明度为0.3，使其不太突兀
    

    

    # --- 修改点 3: 使用正确的列名进行绘图 ---
    ax.plot(df1['relative_time_s'], df1['throughput_tokens_per_sec'], marker='o', linestyle='-', label='Generation Instance 1')
    ax.plot(df2['relative_time_s'], df2['throughput_tokens_per_sec'], marker='s', linestyle='--', label='Generation Instance 2')

    # 设置matplotlib支持中文显示
    # 请确保您的环境中有SimHei字体，或者替换成您有的任何中文字体，如 'Microsoft YaHei'
    try:
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
    except:
        print("警告: 未找到 'SimHei' 字体，图表中的中文可能无法正常显示。")

    ax.set_xlabel('Timeline', fontsize=16)
    # --- 修改点 4: 更新Y轴标签 ---
    ax.set_ylabel('Throughput (tokens/s)', fontsize=16)
    
    # 去除x轴刻度标签
    ax.set_xticklabels([])
    
    # 设置y轴刻度标签字体大小
    ax.tick_params(axis='y', labelsize=15)
    
    ax.legend(fontsize=16, loc='upper center', bbox_to_anchor=(0.5, 1.15), ncol=2, frameon=True)
    
    plt.savefig(output_plot, dpi=300, bbox_inches='tight')
    print(f"图表已成功保存到 '{output_plot}'。")
    plt.close()


if __name__ == '__main__':
    if not os.path.exists(INPUT_FILE_1) or not os.path.exists(INPUT_FILE_2):
        print(f"错误: 无法找到输入文件 {INPUT_FILE_1} 或 {INPUT_FILE_2}")
        print("请确保文件路径在脚本中配置正确。")
    else:
        process_file_to_intermediate(INPUT_FILE_1, INTERMEDIATE_FILE_1)
        
        process_file_to_intermediate(INPUT_FILE_2, INTERMEDIATE_FILE_2)
        
        # 仅在两个中间文件都成功生成后才绘图
        if os.path.exists(INTERMEDIATE_FILE_1) and os.path.exists(INTERMEDIATE_FILE_2):
            label1 = os.path.basename(INPUT_FILE_1)
            label2 = os.path.basename(INPUT_FILE_2)
            plot_from_intermediate_files(INTERMEDIATE_FILE_1, INTERMEDIATE_FILE_2, label1, label2, OUTPUT_PLOT_FILE)
            print("\n所有任务已完成！")
        else:
            print("\n由于一个或多个中间文件未能生成，绘图步骤已跳过。")